/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.analysys.presto.connector.hbase;

import com.facebook.presto.spi.ColumnHandle;
import com.facebook.presto.common.type.Type;
import com.google.common.base.Preconditions;
import io.airlift.log.Logger;
import org.apache.commons.lang.StringUtils;
import org.apache.hadoop.hbase.Cell;
import org.apache.hadoop.hbase.CellUtil;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.Connection;
import org.apache.hadoop.hbase.client.Get;
import org.apache.hadoop.hbase.client.Result;
import org.apache.hadoop.hbase.client.Table;
import org.apache.hadoop.hbase.util.Bytes;
import org.apache.http.client.utils.CloneUtils;

import java.net.InetAddress;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Date;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.stream.Collectors;

import static com.analysys.presto.connector.hbase.utils.Constant.SYSTEMOUT_INTERVAL;
import static java.util.Objects.requireNonNull;

/**
 * Use record.rawCells() api to loop column value, this is 20% faster than result.getValue
 *
 * @author wupeng
 * @date 2019/01/29
 */
public class HBaseGetRecordCursor extends HBaseRecordCursor {


    public static final Logger log = Logger.get(HBaseGetRecordCursor.class);


    /**
     * hbase connection
     */
    private Connection connection;


    /**
     * hbase table
     */
    private Table table;

    /**
     * 结果集索引
     */
    private int currentRecordIndex = 0;


    /**
     * hbase get 查询结果
     */
    private List<Result> results;


    /**
     * 多行
     */
    private final LinkedList<Object[]> rows;


    HBaseGetRecordCursor(List<HBaseColumnHandle> columnHandles,
                         HBaseSplit hbaseSplit,
                         Map<Integer, HBaseColumnHandle> fieldIndexMap,
                         Table table,
                         Connection connection) {
        startTime = System.currentTimeMillis();
        this.columnHandles = columnHandles;
        this.rows = new LinkedList<>();
        this.fieldIndexMap = fieldIndexMap;

        this.rowKeyColName = requireNonNull(hbaseSplit.getRowKeyName(),
                "RowKeyName cannot be null if you want to query by RowKey");

        this.split = hbaseSplit;
        this.connection = connection;
        this.table = table;
        try {
            List<String> rowKeys = hbaseSplit.getConstraint().stream()
                    .map(cond -> (String) cond.getValue()).collect(Collectors.toList());
            this.results = getResults(rowKeys, table);
        } catch (Exception e) {
            log.error(e, e.getMessage());
            this.results = new ArrayList<>(0);
            this.close();
        }
        this.totalBytes = 0L;
    }


    /**
     * hbase get query
     * @param rowKeys
     * @param table
     * @return
     */
    private List<Result> getResults(List<String> rowKeys, Table table) {
        // 查询，字段输出
        long start = System.currentTimeMillis();
        List<Get> gets = rowKeys.stream().map(rowKey -> {
                    Get get = new Get(Bytes.toBytes(rowKey));
                    for (ColumnHandle ch : columnHandles) {
                        HBaseColumnHandle hch = (HBaseColumnHandle) ch;
                        // RowKey column has no column family, so we don't need to do get.addColumn() here.
                        if (this.split.getRowKeyName() != null
                                && this.split.getRowKeyName().equals(hch.getColumnName())) {
                            continue;
                        }

                        // 如果行转列模式，则加载所有列
                        if(!this.split.isExplodeMode()) {
                            get.addColumn(Bytes.toBytes(hch.getFamily()), Bytes.toBytes(hch.getColumnName()));
                        }
                    }
                    return get;
                }
        ).collect(Collectors.toList());
        log.info("get " + rowKeys.size() + " keys result cost time: " + (System.currentTimeMillis() - start) + " ms");
        try {
            return Arrays.stream(table.get(gets)).filter(it->!it.isEmpty()).collect(Collectors.toList());
        } catch (Exception e) {
            log.error(e.getMessage(), e);
            return new ArrayList<>(0);
        }
    }

    @Override
    public boolean advanceNextPosition() {
        try {
            if (this.results != null && this.currentRecordIndex >= this.results.size()) {
                // 没有查询到结果集
                InetAddress localhost = InetAddress.getLocalHost();
                // Random printing
                log.info("BATCH GET RECORD. tableName=" + split.getTableName()
                        + ", rowKey_0=" + split.getConstraint().get(0) + ", READ_DATA_TIME="
                        + (System.currentTimeMillis() - startTime) + " mill secs. recordCount=" + recordCount
                        + ", startTime=" + new Date(startTime).toString() + ", localhost=" + localhost.getHostAddress()
                        + ", specified worker ip: "
                        + (split.getAddresses().size() > 0 ? split.getAddresses().get(0).toString() : ""));
                return false;
            } else {
                Object[] fieldsValue = new Object[this.columnHandles.size()];
                ordinalPositionAndFieldsIndexMap.clear();
                int fieldIndex = 0;
                // 处理一行 hbase result
                Result record = this.results.get(this.currentRecordIndex);
                if(split.isExplodeMode()) {
                    // 需要行转列的模式
                    if(rows.isEmpty()) {
                        long start = System.currentTimeMillis();
                        // 如果没有行了，则扫描下一个
                        List<String> cols = new ArrayList<>();
                        Map<String, byte[]> resultMap = new HashMap<>();
                        for (Cell cell : record.rawCells()) {
                            if(Bytes.toString(CellUtil.cloneQualifier(cell)).startsWith(split.getExplodeColumn() + "_")) {
                                cols.add(Bytes.toString(CellUtil.cloneValue(cell)));
                            }
                            String colName = Bytes.toString(CellUtil.cloneFamily(cell))+":"+Bytes.toString(CellUtil.cloneQualifier(cell));
                            // 字段名称转为小写
                            colName = colName.toLowerCase(Locale.ENGLISH);
                            resultMap.put(colName, CellUtil.cloneValue(cell));
                        }

                        // 从 record 读取所有的 行
                        for (String col : cols) {
                            Object[] objs = new Object[columnHandles.size()];
                            int i = 0;
                            for (HBaseColumnHandle column : columnHandles) {
                                if(column.isIsRowKey()) {
                                    continue;
                                }

                                String colN = column.getFamily()+":"+ column.getColumnName() + "_" + col;
                                // 上面字段名称是转为小写的，这里要转为小写去获取？
                                byte[] r = resultMap.get(StringUtils.lowerCase(colN));
                                if(r == null) {
                                    // 如果 explode column 拼接字段没有获取到，则采用直接的字段拼接来获取
                                    r = resultMap.get(StringUtils.lowerCase(column.getFamily() + ":" + column.getColumnName()));
                                }
                                Object value = null;
                                if(r == null){
                                    value = null;
                                } else if ("BINARY".equalsIgnoreCase(split.getStorageType())) {
                                    this.totalBytes += r.length;
                                    value = matchValueBinary(column.getColumnType(), r);
                                } else {
                                    this.totalBytes += r.length;
                                    value = matchValueText(column.getColumnType(), r);
                                }

                                objs[i] = value;
                                i++;
                            }
                            // 添加一条数据
                            rows.add(objs);
                        }
                        log.info("read hbase explode row has " + rows.size() + " rows, cost time: "+(System.currentTimeMillis()-start) +" ms.");
                    }

                    // 弹出一行数据
                    Object[] row = rows.poll();
                    int i = 0;
                    for (HBaseColumnHandle column : columnHandles) {
                        if(column.isIsRowKey()) {
                            continue;
                        }
                        fieldsValue[fieldIndex] = row[i];
                        ordinalPositionAndFieldsIndexMap.put(column.getOrdinalPosition(), fieldIndex);
                        fieldIndex++;
                        i++;
                    }

                    // Handle the value of rowKey column
                    setRowKeyValue2FieldsAry(record, fieldIndex, fieldsValue);
                    //log.info("process one row cost time: " + (System.currentTimeMillis() - start) + " ms");

                    // 一行读取完成
                    setFieldValues(fieldsValue);
                    if(rows.isEmpty()) {
                        // 行转列完成，下一行开始
                        this.currentRecordIndex++;
                    }
                } else {
                    String colName = null;
                    // 普通模式
                    for (Cell cell : record.rawCells()) {
                        colName = Bytes.toString(CellUtil.cloneQualifier(cell)).toLowerCase(Locale.ENGLISH);
                        HBaseColumnHandle hch = fieldIndexMap.get(colName.hashCode());
                        if (hch == null) {
                            continue;
                        }

                        this.totalBytes += cell.getValueLength();
                        Object value = null;
                        if ("BINARY".equalsIgnoreCase(split.getStorageType())) {
                            value = matchValueBinary(hch.getColumnType(), CellUtil.cloneValue(cell));
                        } else {
                            value = matchValueText(hch.getColumnType(), CellUtil.cloneValue(cell));
                        }
                        fieldsValue[fieldIndex] = value;
                        ordinalPositionAndFieldsIndexMap.put(hch.getOrdinalPosition(), fieldIndex);
                        fieldIndex++;
                    }

                    // Handle the value of rowKey column
                    setRowKeyValue2FieldsAry(record, fieldIndex, fieldsValue);

                    // 一行读取完成
                    setFieldValues(fieldsValue);
                    this.currentRecordIndex++;
                }
                return true;
            }
        } catch (Exception ex) {
            log.error(ex, ex.getMessage());
            this.close();
            fieldIndexMap.forEach((cName, columnHandle) ->
                    log.error("fieldIndexMap: key=" + cName + ", hch.toString=" + columnHandle.toString())
            );
            return false;
        }
    }

    @Override
    public void close() {
        if(table != null) {
            try {
                table.close();
            } catch (Exception e) {
                log.error(e.getMessage(), e);
            }
        }
        if (connection != null) {
            try {
                connection.close();
            } catch (Exception e) {
                log.error(e.getMessage(), e);
            }
        }
    }

    @Override
    public long getCompletedBytes() {
        return this.totalBytes;
    }

    @Override
    public long getReadTimeNanos() {
        return System.currentTimeMillis() - startTime;
    }

    @Override
    public Type getType(int field) {
        Preconditions.checkArgument(field < this.columnHandles.size(), "Invalid field index");
        return this.columnHandles.get(field).getColumnType();
    }
}
